import math
import os
import random
import sys
import bpy
import bpy_extras
from bpy_extras.object_utils import world_to_camera_view
import bmesh
from mathutils import Vector
import glob
sys.path.append(os.path.dirname(__file__))
MAIN_PATH = r"F:\project\lssm_project_all\SIMPAC-2024-292-main\SIMPAC-2024-292-main_new"
DATASET_PATH = r"F:\project\lssm_project_all\SIMPAC-2024-292-main\SIMPAC-2024-292-main_new\dataset"
CAMERA_HEIGHT = 400
render_resolution = [640, 640]
RENDER_SAMPLES = 64
USE_FOLDER = "Cars"
ALL_VEHICLES = ["car", "tractor"]
ADD_TRACTORS = False


# 1.导入地图的两个函数
def get_all_blend_files(lands_folder):
    """收集所有.blend地图文件的绝对路径"""
    all_files = os.listdir(lands_folder)
    blend_files = [f for f in all_files if f.endswith('.blend')]
    blend_files_fullpath = [os.path.join(lands_folder, f) for f in blend_files]
    return blend_files_fullpath


def import_maps_no_repeat(lands_folder):
    """随机打乱，依次导入lands_count个不重复地图"""
    blend_files = get_all_blend_files(lands_folder)
    lands_count = len(blend_files)
    random.shuffle(blend_files)
    selected_blend_files = blend_files[:lands_count]
    return selected_blend_files


def get_ground_z():
    """
    返回地图地面的世界坐标 z。优先用 EXPORT_GOOGLE_SAT_WM 的网格顶点。
    找不到就回退到 0.0。
    """
    wm = bpy.data.objects.get("EXPORT_GOOGLE_SAT_WM")
    if wm and wm.type == 'MESH' and len(wm.data.vertices) > 0:
        M = wm.matrix_world
        # 取前几个顶点求个平均，防止单个顶点有噪声
        n = min(16, len(wm.data.vertices))
        zs = [(M @ wm.data.vertices[i].co).z for i in range(n)]
        return sum(zs) / len(zs)
    return 0.0


# 2.对相机的一些操作
def get_camera_ground_footprint(cam_obj, height):
    """
    返回相机垂直俯视时在 z=0 地面的视场宽高（单位与地图一致）。
    公式：width = 2 * h * tan(FOVx/2), height = 2 * h * tan(FOVy/2)
    """
    camd = cam_obj.data
    # 确保透视相机
    angle_x = camd.angle_x
    angle_y = camd.angle_y
    w = 2.0 * height * math.tan(angle_x * 0.5)
    h = 2.0 * height * math.tan(angle_y * 0.5)
    return w, h


def ensure_camera_looking_down(cam_obj):
    """
    让相机真实“看向”当前相机正下方的地面点（根据地图真实地面 z），
    并设置合理的裁剪面，保证 world_to_camera_view 的 NDC 正常。
    """
    # 设为当前相机
    bpy.context.scene.camera = cam_obj

    # 相机参数统一一下，避免裁剪面导致 ndc.z 异常
    camd = cam_obj.data
    camd.type = 'PERSP'
    camd.clip_start = 0.1
    camd.clip_end = max(5000.0, CAMERA_HEIGHT * 10.0)

    # 计算当前相机正下方的地面目标点
    ground_z = get_ground_z()
    target = Vector((cam_obj.location.x, cam_obj.location.y, ground_z))

    # 让相机本地 -Z 轴对准地面目标点，Y 作为上方向
    direction = (target - cam_obj.location).normalized()
    quat = direction.to_track_quat('-Z', 'Y')

    cam_obj.rotation_mode = 'XYZ'
    cam_obj.rotation_euler = quat.to_euler()

    # 刷新依赖图，然后打印一次 NDC 做校验
    bpy.context.view_layer.update()
    ndc = bpy_extras.object_utils.world_to_camera_view(bpy.context.scene, cam_obj, target)
    print("[DEBUG] center ground NDC:", ndc)  # 正常应接近 (0.5, 0.5, 0..1)


def build_camera_grid_positions(max_x, max_y, cam_obj, margin_ratio=0.05, stride_ratio=0.9):
    """
    用相机地面 footprint 铺网格，筛出“能看到公路”的相机位置。
    判定方式：相机在地面上的矩形视域(宽 fp_w, 高 fp_h)内，是否有任意道路顶点(世界坐标)落入。
    """
    scene = bpy.context.scene
    h = cam_obj.location.z
    fp_w, fp_h = get_camera_ground_footprint(cam_obj, h)

    # 边界留白，避免相机出界
    mx = max_x * (1.0 - margin_ratio)
    my = max_y * (1.0 - margin_ratio)

    # 网格步长
    step_x = fp_w * stride_ratio if fp_w > 0 else max_x / 5
    step_y = fp_h * stride_ratio if fp_h > 0 else max_y / 5

    # 道路顶点 -> 世界坐标
    try:
        road_obj = bpy.data.objects["Ways:highway"]
        M = road_obj.matrix_world
        verts_world = [M @ v.co for v in road_obj.data.vertices]
    except KeyError:
        print("[WARN] 没有 Ways:highway 对象，默认直接铺满网格。")
        verts_world = []

    print("[DEBUG] highway verts:", len(verts_world))
    if verts_world:
        vw0 = verts_world[0]
        print("[DEBUG] sample world vert:", (vw0.x, vw0.y, vw0.z))

    camera_locations = []
    eps = max(min(fp_w, fp_h) * 0.05, 1.0)  # 小余量，避免刚好卡边

    y = -my
    while y <= my + 1e-6:
        x = -mx
        while x <= mx + 1e-6:
            # 相机平移到网格点（朝向已由 ensure_camera_looking_down() 保证为俯视）
            cam_obj.location.x = x
            cam_obj.location.y = y

            # 该相机位置在地面的覆盖矩形
            half_w = fp_w * 0.5
            half_h = fp_h * 0.5
            xmin, xmax = x - half_w - eps, x + half_w + eps
            ymin, ymax = y - half_h - eps, y + half_h + eps

            visible_road = True if not verts_world else False
            if verts_world:
                # 只要有一个道路顶点落入覆盖矩形，就视为可见
                for vw in verts_world:
                    vx, vy = vw.x, vw.y
                    if (xmin <= vx <= xmax) and (ymin <= vy <= ymax):
                        visible_road = True
                        break

            if visible_road:
                camera_locations.append((x, y, h))

            x += step_x
        y += step_y

    return camera_locations


def filter_cameras_inside_map(camera_locations, cam_obj, max_x, max_y):
    """
    过滤：只保留视域 footprint 完全落在地图 [-max_x,max_x] × [-max_y,max_y] 内的相机。
    不改相机的 FOV/高度，只根据 footprint 做几何判定。
    """
    # 用当前相机的高度和FOV算 footprint
    h = cam_obj.location.z
    fp_w, fp_h = get_camera_ground_footprint(cam_obj, h)
    half_w = fp_w * 0.5
    half_h = fp_h * 0.5

    kept = []
    for (x, y, z) in camera_locations:
        xmin, xmax = x - half_w, x + half_w
        ymin, ymax = y - half_h, y + half_h
        fully_inside = (xmin >= -max_x) and (xmax <= max_x) and (ymin >= -max_y) and (ymax <= max_y)
        if fully_inside:
            kept.append((x, y, z))
    print(f"[INFO] inside-map cameras: {len(kept)} / {len(camera_locations)}")
    return kept


# 3.添加太阳
def ensure_single_sun(name="KeySun", energy=5.0, elevation_deg=45.0, azimuth_deg=30.0, soft_size=0.2):
    # 1) 找现有主灯
    sun = bpy.data.objects.get(name)
    # 2) 删除其他多余的 SUN 灯
    for obj in list(bpy.data.objects):
        if obj.type == 'LIGHT' and obj.data.type == 'SUN' and obj.name != name:
            bpy.data.objects.remove(obj, do_unlink=True)

    # 3) 没有就创建
    if sun is None:
        bpy.ops.object.light_add(type="SUN")
        sun = bpy.context.object
        sun.name = name

    # 4) 统一参数
    sun.data.energy = energy
    sun.data.use_shadow = True
    sun.data.shadow_soft_size = soft_size
    sun.rotation_euler = (math.radians(elevation_deg), 0.0, math.radians(azimuth_deg))

    # 5) 世界背景强度（可固定一次，不要每帧随机）
    world = bpy.data.worlds.get("World")
    if world and world.node_tree:
        bg = world.node_tree.nodes.get("Background")
        if bg:
            bg.inputs[1].default_value = 0.9  # 固定，不会越叠越亮


# 4.删除所有的车并添加新车
def delete_all_cars():
    roots = [o for o in bpy.data.objects if o.name.startswith("Car_")]
    for r in roots:
        # 先删子，再删根
        for ch in list(r.children_recursive):
            bpy.data.objects.remove(ch, do_unlink=True)
        bpy.data.objects.remove(r, do_unlink=True)


def sample_random_road_point_in_view(cam_obj, max_trials=200):
    """
    在相机视野内，随机选 Ways:highway 的一条边，在边上取一点 (x,y,z) 和道路朝向 angle。
    返回: ((x,y,z), angle) 或 None
    """
    scene = bpy.context.scene

    # 进入编辑态以便拿到 BMesh 边
    try:
        bpy.ops.object.select_all(action="DESELECT")
        road = bpy.data.objects["Ways:highway"]
    except KeyError:
        print("[ERROR] 没有 Ways:highway 对象")
        return None

    bpy.context.view_layer.objects.active = road
    bpy.ops.object.mode_set(mode="EDIT")
    bm = bmesh.from_edit_mesh(road.data)
    bm.edges.ensure_lookup_table()

    trials = 0
    while trials < max_trials:
        trials += 1
        # 随机取一条边
        e = bm.edges[random.randint(0, len(bm.edges) - 1)]
        v1, v2 = e.verts[0], e.verts[1]
        # 转世界坐标
        M = road.matrix_world
        p1 = M @ v1.co
        p2 = M @ v2.co

        # 沿边随机取一点 t∈[0,1]
        t = random.random()
        x = (1 - t) * p1.x + t * p2.x
        y = (1 - t) * p1.y + t * p2.y
        z = (1 - t) * p1.z + t * p2.z

        # 判断这个地面点是否在相机前方且落在成像 0..1 区域
        if not is_world_point_visible_on_ground(scene, cam_obj, (x, y)):
            continue

        # 朝向 = 边的方向角（XY 平面）
        dx = (p2.x - p1.x)
        dy = (p2.y - p1.y)
        if abs(dx) + abs(dy) < 1e-6:
            continue
        angle = math.atan2(dy, dx)

        # 地面高度（贴地更稳）
        gz = ground_z_at(x, y, z_hint=cam_obj.location.z)
        return ((float(x), float(y), float(gz)), float(angle))

    print("[WARN] 视野内未采到道路点")
    bpy.ops.object.mode_set(mode="OBJECT")
    return None


def get_object_xy_aabb(obj):
    """返回对象(含子Mesh)在世界坐标地面上的XY轴对齐包围框 (minx,miny,maxx,maxy)。"""
    minx = miny = float('inf')
    maxx = maxy = float('-inf')
    has_pts = False
    for ch in obj.children_recursive:
        if ch.type == "MESH":
            for c in ch.bound_box:
                w = ch.matrix_world @ Vector(c)
                minx = min(minx, w.x); maxx = max(maxx, w.x)
                miny = min(miny, w.y); maxy = max(maxy, w.y)
                has_pts = True
    if obj.type == "MESH":  # 根本身就是Mesh的情况
        for c in obj.bound_box:
            w = obj.matrix_world @ Vector(c)
            minx = min(minx, w.x); maxx = max(maxx, w.x)
            miny = min(miny, w.y); maxy = max(maxy, w.y)
            has_pts = True
    if not has_pts:
        return None
    return (minx, miny, maxx, maxy)


def rects_overlap(a, b, margin_m=0.6):
    """两个XY包围框是否重叠；margin_m为安全边距（米/Blender单位）。"""
    (ax1, ay1, ax2, ay2) = (a[0]-margin_m, a[1]-margin_m, a[2]+margin_m, a[3]+margin_m)
    (bx1, by1, bx2, by2) = (b[0]-margin_m, b[1]-margin_m, b[2]+margin_m, b[3]+margin_m)
    # 分离轴：只要有一轴不相交就“非重叠”
    if ax2 < bx1 or bx2 < ax1:
        return False
    if ay2 < by1 or by2 < ay1:
        return False
    return True


def add_new_car(cx, cy, use_folder, heading=None):
    """
    导入一辆车并放在 (cx, cy) 正下方：
    - 收编本次导入的所有 FBX 子对象到一个 Car_XXXXX 根（Empty）
    - 先把 root 放到“导入后几何中心”，再随机旋转（绕几何中心）
    - 之后贴地，再把整体几何中心对齐到 (cx,cy)
    - 返回 root（失败返回 None）
    """
    cars_folder = os.path.join(MAIN_PATH, "Vehicles", use_folder)
    car_model_files = glob.glob(os.path.join(cars_folder, "*.fbx"))
    if not car_model_files:
        print("No car models found in", cars_folder)
        return None

    before = set(bpy.data.objects)

    car_model_path = random.choice(car_model_files)
    print(f"Selected car FBX: {car_model_path}", flush=True)
    bpy.ops.import_scene.fbx(filepath=car_model_path)

    after = set(bpy.data.objects)
    new_objs = [o for o in (after - before)]
    if not new_objs:
        print("[ERROR] 导入后没有新对象")
        return None

    # 创建根 Empty，并把新导入对象全挂到根下（保持世界矩阵不变）
    unique_id = f"{random.randint(0, 99999):05d}"
    root = bpy.data.objects.new(f"Car_{unique_id}", None)
    bpy.context.scene.collection.objects.link(root)
    for o in new_objs:
        try:
            M = o.matrix_world.copy()
            o.parent = root
            o.matrix_parent_inverse = root.matrix_world.inverted()
            o.matrix_world = M
        except Exception as e:
            print("[WARN] 绑定子对象失败:", o.name, e)

    # —— 计算“导入后”的整体几何中心（用于让旋转绕中心进行）——
    all_world_corners = []
    for o in root.children_recursive:
        if o.type == "MESH":
            all_world_corners.extend([o.matrix_world @ Vector(c) for c in o.bound_box])
    if not all_world_corners:
        print("[ERROR] 新导入对象没有 Mesh")
        return None
    cx0 = sum(v.x for v in all_world_corners) / len(all_world_corners)
    cy0 = sum(v.y for v in all_world_corners) / len(all_world_corners)
    cz0 = sum(v.z for v in all_world_corners) / len(all_world_corners)

    # 1) 先把 root 放到“导入后几何中心”，确保随机旋转绕几何中心
    root.location = (cx0, cy0, cz0)

    # 2) 再随机朝向（绕中心旋转，不会把车甩飞）
    root.rotation_mode = "XYZ"
    if heading is None:
        angle = random.uniform(0, math.pi * 2)
    else:
        angle = heading
    root.rotation_euler = (0.0, 0.0, angle)

    # 3) 旋转后重新算整体最低点，做“贴地”
    all_world_corners = []
    for o in root.children_recursive:
        if o.type == "MESH":
            all_world_corners.extend([o.matrix_world @ Vector(c) for c in o.bound_box])
    min_z = min(v.z for v in all_world_corners)
    z_offset = -min_z
    gz = ground_z_at(cx, cy, CAMERA_HEIGHT)
    root.location = (cx, cy, gz + z_offset + 0.02)

    # 6) 兜底：确保可渲染 & 链接根到场景集合
    try:
        if bpy.context.scene.collection not in (root.users_collection or []):
            try:
                bpy.context.scene.collection.objects.link(root)
            except RuntimeError:
                pass
        ensure_renderable_object_tree(root)
    except Exception as e:
        print("[WARN] ensure_renderable_object_tree failed:", e)

    bpy.context.view_layer.update()
    return root


def is_object_fully_in_frame(obj, cam_obj, margin=0.02):
    """
    检查 obj（含所有子mesh）的投影是否完全落在相机视野内，留 margin 边距。
    返回 (ok, bbox_ndc)；bbox_ndc=(minx,miny,maxx,maxy) 供调试。
    """
    scene = bpy.context.scene
    pts_ndc = []

    for ch in obj.children_recursive:
        if ch.type == "MESH":
            for c in ch.bound_box:  # 8个角
                w = ch.matrix_world @ Vector(c)
                ndc = bpy_extras.object_utils.world_to_camera_view(scene, cam_obj, w)
                pts_ndc.append(ndc)

    if not pts_ndc:
        return False, None

    minx = min(p.x for p in pts_ndc)
    miny = min(p.y for p in pts_ndc)
    maxx = max(p.x for p in pts_ndc)
    maxy = max(p.y for p in pts_ndc)
    minz = min(p.z for p in pts_ndc)

    ok = (
        minz >= 0.0 and
        minx >= 0.0 + margin and maxx <= 1.0 - margin and
        miny >= 0.0 + margin and maxy <= 1.0 - margin
    )
    return ok, (minx, miny, maxx, maxy)


def add_shadow_catcher():
    """
    Adds a plane that will be used as a shadow catcher - to
    receive and show the shadow while being transparent.
    """
    sc = bpy.data.objects.get("Plane")
    if sc is None:
        bpy.ops.mesh.primitive_plane_add(size=1, location=(0, 0, 0.1))
        sc = bpy.context.active_object
        sc.name = "Plane"
        sc.scale = [
            bpy.data.objects["EXPORT_GOOGLE_SAT_WM"].dimensions.x,
            bpy.data.objects["EXPORT_GOOGLE_SAT_WM"].dimensions.y,
            1]
        sc.is_shadow_catcher = True
        sc.visible_diffuse = False
        sc.visible_glossy = False
        bpy.context.scene.render.film_transparent = True


def prepare_compositor():
    """
    初始化合成器节点：BG_IMAGE_NODE（背景图） + ALPHA_OVER（前景叠加）。
    若已存在，则不重复创建。
    前景链路（Render Layers）可按你原有的 hue/exposure/blur/scale/pixelate 继续串。
    """
    scene = bpy.context.scene
    scene.use_nodes = True
    tree = scene.node_tree

    # 若已配置过，直接返回
    if tree.nodes.get("BG_IMAGE_NODE") and tree.nodes.get("ALPHA_OVER"):
        return

    # 基础节点
    rl = tree.nodes.get("Render Layers") or tree.nodes.new("CompositorNodeRLayers")
    comp = tree.nodes.get("Composite") or tree.nodes.new("CompositorNodeComposite")
    viewer = tree.nodes.get("Viewer") or tree.nodes.new("CompositorNodeViewer")

    # 背景图节点
    img = tree.nodes.new("CompositorNodeImage")
    img.name = "BG_IMAGE_NODE"

    # 颜色/曝光（你的原链路）
    hue = tree.nodes.new("CompositorNodeHueSat")
    exp = tree.nodes.new("CompositorNodeExposure")
    hue.inputs[2].default_value = random.uniform(0.9, 1.0)
    exp.inputs[1].default_value = 0.0  # 替换 random.uniform(1.0, 1.3)
    # exp.inputs[1].default_value = random.uniform(1.0, 1.3)

    # 模糊/缩放/像素化（按你原参数）
    blur1 = tree.nodes.new("CompositorNodeBlur"); blur1.filter_type = "CATROM"
    blur2 = tree.nodes.new("CompositorNodeBlur"); blur2.filter_type = "CATROM"
    blur1.size_x = round(render_resolution[0] / 341)
    blur1.size_y = round(render_resolution[1] / 341)
    blur2.size_x = round(render_resolution[0] / 1024)
    blur2.size_y = round(render_resolution[1] / 1024)

    scale1 = tree.nodes.new("CompositorNodeScale")
    scale2 = tree.nodes.new("CompositorNodeScale")
    scale1.inputs[1].default_value = 0.9
    scale1.inputs[2].default_value = 0.9
    scale2.inputs[1].default_value = 1.0 / scale1.inputs[1].default_value
    scale2.inputs[2].default_value = 1.0 / scale1.inputs[2].default_value

    pixel = tree.nodes.new("CompositorNodePixelate")

    # 叠加
    over = tree.nodes.new("CompositorNodeAlphaOver")
    over.name = "ALPHA_OVER"
    over.inputs[0].default_value = 1.0  # 混合因子

    # 连接：背景 -> over[1]；前景(渲层处理后) -> over[2]；over -> 输出
    links = tree.links
    links.new(img.outputs[0], over.inputs[1])
    links.new(rl.outputs[0], hue.inputs[0])
    links.new(hue.outputs[0], exp.inputs[0])
    links.new(exp.outputs[0], blur1.inputs[0])
    links.new(blur1.outputs[0], scale1.inputs[0])
    links.new(scale1.outputs[0], pixel.inputs[0])
    links.new(pixel.outputs[0], scale2.inputs[0])
    links.new(scale2.outputs[0], blur2.inputs[0])
    links.new(blur2.outputs[0], over.inputs[2])
    links.new(over.outputs[0], comp.inputs[0])
    links.new(over.outputs[0], viewer.inputs[0])


def edit_compositor(background_filepath, background_filename):
    """
    将合成器中的 BG_IMAGE_NODE 指向本帧的 backgrounds/{idx}.png
    """
    tree = bpy.context.scene.node_tree
    img = tree.nodes.get("BG_IMAGE_NODE")
    if img is None:
        raise RuntimeError("合成器未初始化：缺少 BG_IMAGE_NODE。先调用 prepare_compositor().")
    bpy.data.images.load(background_filepath, check_existing=True)
    img.image = bpy.data.images[background_filename]


def render_backgrounds(camera_location, curr_index_global):
    """
    渲一张纯背景到 backgrounds/{idx}.png：
    - 临时禁用合成器
    - 临时隐藏车辆与 shadow catcher
    - 只保留地图本体
    """
    scene = bpy.context.scene
    cam = bpy.data.objects["Camera"]

    prev_use_nodes = scene.use_nodes
    scene.use_nodes = False
    _toggle_cars(True)
    _toggle_shadow_catcher(True)

    try:
        scene.render.engine = "CYCLES"
        scene.cycles.device = "GPU"
        scene.cycles.samples = RENDER_SAMPLES
        scene.render.resolution_x, scene.render.resolution_y = render_resolution

        # 定位相机
        if isinstance(camera_location, tuple):
            cam.location = camera_location
        else:
            cam.location = tuple(camera_location)

        scene.camera = cam
        bpy.context.view_layer.update()

        bg_name = f"{curr_index_global}.png"
        bg_path = os.path.join(DATASET_PATH, "backgrounds", bg_name)
        scene.render.filepath = bg_path
        bpy.ops.render.render(write_still=True)
    finally:
        # 恢复场景状态
        scene.use_nodes = prev_use_nodes
        _toggle_shadow_catcher(False)
        _toggle_cars(False)


def annotate_vehicles_yolo(all_vehicle_boxes, scene, cam_obj, curr_index, vehicle_type):
    label_path = os.path.join(DATASET_PATH, "labels", f"{curr_index}.txt")
    with open(label_path, "a", encoding="utf-8") as f:
        for pts_world in all_vehicle_boxes:
            xs, ys = [], []
            for p in pts_world:
                ndc = world_to_camera_view(scene, cam_obj, p)
                if ndc.z < 0.0:  # 在相机后方
                    continue
                # —— 关键：把 Blender 的底原点改成 YOLO 的顶原点 —— #
                x = float(ndc.x)
                y = 1.0 - float(ndc.y)

                # clamp 到 [0,1]
                x = 0.0 if x < 0.0 else 1.0 if x > 1.0 else x
                y = 0.0 if y < 0.0 else 1.0 if y > 1.0 else y

                xs.append(x)
                ys.append(y)

            if not xs or not ys:
                continue

            bbox = return_yolo_box(xs, ys)
            if bbox is None:
                continue

            category_id = ALL_VEHICLES.index(vehicle_type)
            cx, cy, w, h = bbox
            f.write(f"{category_id} {cx} {cy} {w} {h}\n")


def is_world_point_visible_on_ground(scene, cam_obj, xy):
    """
    判断地面点 (x, y, 0) 是否在相机成像范围内（0..1）且在相机前方。
    """
    co = Vector((xy[0], xy[1], 0.0))
    co_ndc = bpy_extras.object_utils.world_to_camera_view(scene, cam_obj, co)
    return (0.0 <= co_ndc.x <= 1.0) and (0.0 <= co_ndc.y <= 1.0) and (co_ndc.z >= 0.0)


def render_single_image_with_annotations( cam_obj, curr_index, all_tractor_boxes=None, vehicle_type_car="car"):
    """
    只渲染当前相机位置的一张图（合成+标注）。
    依赖：背景 PNG 已经提前渲染好；已调用 prepare_compositor()，且已 add_shadow_catcher()。
    """
    scene = bpy.context.scene

    # 切换本张背景
    image_filename = f"{curr_index}.png"
    background_filepath = os.path.join(DATASET_PATH, "backgrounds", image_filename)
    edit_compositor(background_filepath, image_filename)

    # 计算车辆 3D 包围盒（你已有的函数）
    all_car_boxes = get_all_cars_bounding_boxes()

    # 渲染出图
    image_filepath = os.path.join(DATASET_PATH, "images", image_filename)
    scene.render.filepath = image_filepath
    bpy.ops.render.render(write_still=True)

    # 清空 YOLO 标签文件
    with open(os.path.join(DATASET_PATH, "labels", f"{curr_index}.txt"), "w", encoding="utf-8"):
        pass

    # 写 YOLO 标注
    annotate_vehicles_yolo(all_car_boxes, scene, cam_obj, curr_index, vehicle_type_car)
    if ADD_TRACTORS and all_tractor_boxes:
        annotate_vehicles_yolo(all_tractor_boxes, scene, cam_obj, curr_index, "tractor")


def ground_z_at(x, y, z_hint=None):
    """
    返回 (x,y) 正下方地图地面的 z。用 EXPORT_GOOGLE_SAT_WM 做一次从上到下的 ray_cast。
    找不到则回退到 get_ground_z() 或 0.0。
    """
    wm = bpy.data.objects.get("EXPORT_GOOGLE_SAT_WM")
    if not wm or wm.type != 'MESH' or not wm.data.vertices:
        return get_ground_z()
    if z_hint is None:
        z_hint = CAMERA_HEIGHT
    M = wm.matrix_world
    Minv = M.inverted()
    origin_world = Vector((float(x), float(y), float(z_hint)))
    dir_world = Vector((0.0, 0.0, -1.0))
    origin_local = Minv @ origin_world
    dir_local = (Minv.to_3x3() @ dir_world).normalized()
    ok, hit_loc, _, _ = wm.ray_cast(origin_local, dir_local)
    if ok:
        hit_world = M @ hit_loc
        return float(hit_world.z)
    return get_ground_z()


def ensure_renderable_object_tree(root_obj):
    """Make sure object, its children, and its collections/view layer are renderable."""
    def _walk(o):
        try: o.hide_render = False
        except: pass
        if hasattr(o, "hide_viewport"):
            try: o.hide_viewport = False
            except: pass
        for ch in o.children: _walk(ch)
    _walk(root_obj)

    for col in getattr(root_obj, "users_collection", []) or []:
        if hasattr(col, "hide_render"):
            try: col.hide_render = False
            except: pass
        if hasattr(col, "hide_viewport"):
            try: col.hide_viewport = False
            except: pass

    layer = bpy.context.view_layer
    target_cols = set(getattr(root_obj, "users_collection", []) or [])
    def _walk_layer(layer_collection):
        c = layer_collection.collection
        if c in target_cols:
            if hasattr(layer_collection, "exclude"): layer_collection.exclude = False
            if hasattr(layer_collection, "holdout"): layer_collection.holdout = False
            if hasattr(layer_collection, "indirect_only"): layer_collection.indirect_only = False
        for child in layer_collection.children:
            _walk_layer(child)
    _walk_layer(layer.layer_collection)


def _toggle_cars(hide: bool):
    car_roots = [o for o in bpy.data.objects if o.name.startswith("Car_")]
    def _walk(o):
        o.hide_render = hide
        for ch in o.children:
            _walk(ch)
    for r in car_roots:
        _walk(r)


def _get_shadow_catcher():
    return bpy.data.objects.get("Plane")  # 你创建时就叫 Plane


def _toggle_shadow_catcher(hide: bool):
    sc = _get_shadow_catcher()
    if sc:
        sc.hide_render = hide


def get_all_cars_bounding_boxes():
    """返回每辆车的一组世界坐标点（汇总其所有子Mesh的 bound_box 角点）。
       annotate_vehicles_yolo 会把这些点投影成2D再取min/max。"""
    boxes = []
    # 找到每辆车的根（Empty 或 Mesh），名字以 Car_ 开头
    car_roots = [o for o in bpy.data.objects if o.name.startswith("Car_") and o.type in {"EMPTY","MESH"}]
    for root in car_roots:
        pts = []
        # 递归收集所有子Mesh的8个包围盒角点（变换到世界坐标）
        for ch in root.children_recursive:
            if ch.type == "MESH":
                for c in ch.bound_box:
                    pts.append(ch.matrix_world @ Vector(c))
        # 有些FBX根本身就是Mesh，连它自己的角点也收一下
        if root.type == "MESH":
            for c in root.bound_box:
                pts.append(root.matrix_world @ Vector(c))
        if pts:
            boxes.append(pts)
    return boxes


def return_yolo_box(xs, ys):
    """xs/ys 为投影后的 0..1 坐标列表。返回 (cx, cy, w, h) 或 None(过小/异常)。"""
    if not xs or not ys:
        return None
    x1, x2 = min(xs), max(xs)
    y1, y2 = min(ys), max(ys)
    # 规整到 [0,1]
    x1 = max(0.0, min(1.0, x1)); x2 = max(0.0, min(1.0, x2))
    y1 = max(0.0, min(1.0, y1)); y2 = max(0.0, min(1.0, y2))
    w = x2 - x1
    h = y2 - y1
    # 用像素做最小阈值（避免极小物体被写成 0 宽高）
    min_w = 1.0 / float(render_resolution[0])
    min_h = 1.0 / float(render_resolution[1])
    if w < min_w or h < min_h:
        return None
    cx = (x1 + x2) * 0.5
    cy = (y1 + y2) * 0.5
    return (float(cx), float(cy), float(w), float(h))


def main():
    for sub in ["labels", "images", "backgrounds"]:
        d = os.path.join(DATASET_PATH, sub)
        os.makedirs(d, exist_ok=True)

    bpy.context.scene.view_settings.view_transform = "Standard"

    # 1.导入所有地图
    lands_folder = os.path.join(MAIN_PATH, "Lands")
    selected_blend_files = import_maps_no_repeat(lands_folder)
    curr_index_global = 0   # 生成结果的索引

    for blend_path in selected_blend_files:
        print("Importing:", blend_path)
        bpy.ops.wm.open_mainfile(filepath=blend_path)

        # 地图尺寸
        max_x = bpy.data.objects["EXPORT_GOOGLE_SAT_WM"].dimensions.x / 2
        max_y = bpy.data.objects["EXPORT_GOOGLE_SAT_WM"].dimensions.y / 2

        # 保证场上只有一个相机
        for obj in list(bpy.data.objects):
            if obj.type == "CAMERA":
                bpy.data.objects.remove(obj, do_unlink=True)

        bpy.ops.object.camera_add(location=(0, 0, CAMERA_HEIGHT))
        cam = bpy.context.active_object
        cam.name = "Camera"
        bpy.context.scene.camera = cam
        ensure_camera_looking_down(cam)

        # 基于 footprint 自动生成可用相机位置（能看到道路 & 不出界）
        camera_locations_all = build_camera_grid_positions(max_x, max_y, cam)
        camera_locations_all = filter_cameras_inside_map(camera_locations_all, cam, max_x, max_y)

        if not camera_locations_all:
            print("[WARN] 本地图无可用相机位置，跳过。")
            continue

        print(f"[INFO] 可容纳相机位置总数: {len(camera_locations_all)}")
        ensure_single_sun()  # 给世界添加太阳
        for loc in camera_locations_all:
            delete_all_cars()
            # 定位相机并俯视
            cam.location = loc
            ensure_camera_looking_down(cam)
            placed = False
            placed_rects = []   # 记录本帧已放置车辆的XY包围框
            MAX_TRY_PLACE = 60  # 给足尝试次数以避免边缘点
            for attempt in range(MAX_TRY_PLACE):
                # 先在视野内采一个道路点+朝向
                sample = sample_random_road_point_in_view(cam)
                if sample is None:
                    # 退化：相机中心（不推荐，但保证不空）
                    gz = ground_z_at(loc[0], loc[1], z_hint=cam.location.z)
                    road_xyz = (loc[0], loc[1], gz)
                    road_angle = random.uniform(0, 2 * math.pi)
                else:
                    (road_xyz, road_angle) = sample

                # 放车
                car = add_new_car(road_xyz[0], road_xyz[1], use_folder=USE_FOLDER, heading=road_angle)
                if not car:
                    continue

                # 出画检查
                ok, bb = is_object_fully_in_frame(car, cam, margin=0.02)
                if not ok:
                    print(f"[RETRY] out-of-frame, resample... {attempt + 1}")
                    for o in list(car.children_recursive) + [car]:
                        try:
                            bpy.data.objects.remove(o, do_unlink=True)
                        except:
                            pass
                    continue
                # 重叠检查
                aabb = get_object_xy_aabb(car)
                overlapped = any(rects_overlap(aabb, r, margin_m=0.6) for r in placed_rects) if aabb else False
                if overlapped:
                    print(f"[RETRY] overlap with previous car, resample... {attempt + 1}")
                    for o in list(car.children_recursive) + [car]:
                        try:
                            bpy.data.objects.remove(o, do_unlink=True)
                        except:
                            pass
                    continue
                if aabb:
                    placed_rects.append(aabb)
                placed = True
                break
            if not placed:
                print("[WARN] 多次尝试仍无法将车完全放入画面，跳过此相机位置。")
                curr_index_global += 1
                continue
            render_backgrounds(loc, curr_index_global)

            add_shadow_catcher()
            prepare_compositor()

            render_single_image_with_annotations(cam, curr_index_global, all_tractor_boxes=None, vehicle_type_car="car"
            )
            curr_index_global += 1


if __name__ == "__main__":
    main()
